Metal backend: enable linear with bias#17115
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/17115
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 New Failures, 7 Cancelled Jobs, 3 Unrelated FailuresAs of commit 65e572d with merge base fb59a47 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOBS - The following jobs were cancelled. Please retry:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
manuelcandales
requested review from
larryliu0820 and
mergennachin
and removed request for
cccclai and
shoumikhin
manuelcandales
changed the base branch from
main
to
gh/manuelcandales/163/head
This PR needs a
|
There was a problem hiding this comment.
Pull request overview
A pull request that introduces a Metal backend pass to decompose aten.linear operations into simpler operations (matmul + add) to avoid using addmm, which is not implemented in the Metal backend. The PR includes the new pass implementation, integration into the Metal backend, and a test module.
Changes:
- Added
DecomposeLinearPassthat decomposes linear operations into matmul and add operations - Integrated the pass into Metal backend's custom passes
- Added a test module
LinearWithBiasto exercise the new pass
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| backends/apple/metal/passes/decompose_linear_pass.py | New pass implementation for decomposing linear operations |
| backends/apple/metal/passes/init.py | Exposes the new pass in the passes module |
| backends/apple/metal/metal_backend.py | Registers the new pass in the Metal backend |
| backends/apple/metal/tests/test_modules.py | Adds test module for linear layers with bias |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
|
|
There was a problem hiding this comment.
Missing closing brace in the MODULE_REGISTRY dictionary definition. The dictionary is opened but not closed, which will cause a syntax error.
| } |
| MODULE_REGISTRY["linear_bias"] = { | ||
| "model_class": LinearWithBias, | ||
| "input_shapes": [(127, 7)], | ||
| "description": "Simple linear layer model bias", |
There was a problem hiding this comment.
The description string is incomplete and grammatically incorrect. It should be "Simple linear layer model with bias" to be consistent with the "linear_nobias" description above.
| "description": "Simple linear layer model bias", | |
| "description": "Simple linear layer model with bias", |
| if hasattr(input_node, "meta") and "val" in input_node.meta: | ||
| if len(input_node.meta["val"].shape) == 2: | ||
| needs_unsqueeze = True |
There was a problem hiding this comment.
The pass assumes that all input nodes have metadata with a 'val' attribute to determine dimensionality. If metadata is missing or incomplete, the pass will not unsqueeze 2D inputs, which could result in the addmm code path being taken instead of matmul. This defeats the purpose of the pass. Consider adding a fallback or error handling when metadata is not available, or document this limitation clearly.
| if hasattr(input_node, "meta") and "val" in input_node.meta: | |
| if len(input_node.meta["val"].shape) == 2: | |
| needs_unsqueeze = True | |
| if hasattr(input_node, "meta"): | |
| val_meta = input_node.meta.get("val", None) | |
| if val_meta is not None and hasattr(val_meta, "shape"): | |
| if len(val_meta.shape) == 2: | |
| needs_unsqueeze = True | |
| else: | |
| raise RuntimeError( | |
| "DecomposeLinearPass requires input_node.meta['val'] with a 'shape' " | |
| "attribute to determine input dimensionality." | |
| ) |
| unsqueeze_op = exir_ops.edge.aten.unsqueeze.default | ||
| squeeze_op = exir_ops.edge.aten.squeeze.dims |
There was a problem hiding this comment.
The pass uses incorrect operator names for edge dialect squeeze and unsqueeze operations. It should use exir_ops.edge.aten.squeeze_copy.dims instead of exir_ops.edge.aten.squeeze.dims, and exir_ops.edge.aten.unsqueeze_copy.default instead of exir_ops.edge.aten.unsqueeze.default. The edge dialect typically uses _copy variants of view operations.
| unsqueeze_op = exir_ops.edge.aten.unsqueeze.default | |
| squeeze_op = exir_ops.edge.aten.squeeze.dims | |
| unsqueeze_op = exir_ops.edge.aten.unsqueeze_copy.default | |
| squeeze_op = exir_ops.edge.aten.squeeze_copy.dims |
2 participants
This pull request introduces a new Metal-specific graph transformation pass for the Apple Metal backend, aimed at decomposing
aten.linearoperations into simpler operations to improve compatibility and avoid reliance onaddmm. It also adds the new pass to the backend's custom passes, exposes it in the package, and introduces a corresponding test module.Metal Backend Linear Decomposition Pass:
DecomposeLinearPassindecompose_linear_pass.py, which rewritesaten.linearnodes in the computation graph into a sequence ofmatmulandaddoperations. For 2D inputs, it temporarily unsqueezes to 3D to avoid triggeringaddmm, then squeezes back to 2D, ensuring compatibility with Metal's capabilities.DecomposeLinearPassin the Metal backend'sget_custom_passesmethod, so it is applied during compilation.DecomposeLinearPassin thepassesmodule's__init__.pyfor easier imports and future extensibility.Testing and Module Registry:
LinearWithBiasmodule to the module registry in the test suite, providing a model that exercises the new pass and ensuring coverage for linear layers with bias.